
/**
  ******************************************************************************
  * Copyright 2021 The Microbee Authors. All Rights Reserved.
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  * 
  * http://www.apache.org/licenses/LICENSE-2.0
  * 
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  * 
  * @file       AccelCalibrator.c
  * @author     baiyang
  * @date       2022-3-13
  ******************************************************************************
  */

/*----------------------------------include-----------------------------------*/
#include "AccelCalibrator.h"

#include <string.h>
#include <float.h>
#include <rtthread.h>

#include <common/time/gp_time.h>
#include <common/gp_math/gp_matrix_alg.h>
/*-----------------------------------macro------------------------------------*/

/*----------------------------------typedef-----------------------------------*/

/*---------------------------------prototype----------------------------------*/
static void accelcal_set_status(accelcal_t accelcal, enum accel_cal_status_t status);
static bool accelcal_accept_sample(accelcal_t accelcal, const Vector3f_t* sample);
static void accelcal_run_fit(accelcal_t accelcal, uint8_t max_iterations, float* fitness);
static float accelcal_calc_mean_squared_residuals(accelcal_t accelcal, const struct accel_cal_param_t* params);
static float accelcal_calc_residual(const Vector3f_t* sample, const struct accel_cal_param_t* params);
static void accelcal_calc_jacob(accelcal_t accelcal, const Vector3f_t* sample, const struct accel_cal_param_t* params, VectorP ret);
static uint8_t accelcal_get_num_params(accelcal_t accelcal);
static bool accelcal_accept_result(accelcal_t accelcal);
/*----------------------------------variable----------------------------------*/

/*-------------------------------------os-------------------------------------*/

/*----------------------------------function----------------------------------*/
/**
  * @brief       
  * @param[in]   accelcal  
  * @param[out]  
  * @retval      
  * @note        
  */
void accelcal_ctor(accelcal_t accelcal)
{
    rt_memset(accelcal, 0, sizeof(struct accel_calibrator));

    accelcal->_conf_tolerance = ACCEL_CAL_TOLERANCE;
    accelcal->_sample_buffer = NULL;
}

// set Accel calibrator status to make itself ready for future accel cals
void accelcal_clear(accelcal_t accelcal)
{
    accelcal_set_status(accelcal, ACCEL_CAL_NOT_STARTED);
}

/*
    Select options, initialise variables and initiate accel calibration
    Options:
    Fit Type:       Will assume that if accelerometer static samples around all possible orientatio
                    are spread in space will cover a surface of AXIS_ALIGNED_ELLIPSOID or any general 
                    ELLIPSOID as chosen

    Num Samples:    Number of samples user should will be gathering, please note that with type of surface
                    chosen the minimum number of samples required will vary, for instance in the case of AXIS
                    ALIGNED ELLIPSOID atleast 6 distinct samples are required while for generic ELLIPSOIDAL fit
                    which will include calculation of offdiagonal parameters too requires atleast 8 parameters
                    to form a distinct ELLIPSOID

    Sample Time:    Time over which the samples will be gathered and averaged to add to a single sample for least
                    square regression

    offset,diag,offdiag: initial parameter values for LSQ calculation
*/
//default,fit_type = ACCEL_CAL_AXIS_ALIGNED_ELLIPSOID, num_samples = 6, sample_time = 0.5f
void accelcal_start(accelcal_t accelcal, enum accel_cal_fit_type_t fit_type, uint8_t num_samples, float sample_time) {
    Vector3f_t offset = {0.0f, 0.0f, 0.0f};
    Vector3f_t diag = {1.0f, 1.0f, 1.0f};
    Vector3f_t offdiag = {0.0f, 0.0f, 0.0f};
    accelcal_start2(accelcal, fit_type, num_samples, sample_time, &offset, &diag, &offdiag);
}

void accelcal_start2(accelcal_t accelcal, enum accel_cal_fit_type_t fit_type, uint8_t num_samples, float sample_time, Vector3f_t* offset, Vector3f_t* diag, Vector3f_t* offdiag) {
    if (accelcal->_status == ACCEL_CAL_FAILED || accelcal->_status == ACCEL_CAL_SUCCESS) {
        accelcal_clear(accelcal);
    }
    if (accelcal->_status != ACCEL_CAL_NOT_STARTED) {
        return;
    }

    accelcal->_conf_num_samples = num_samples;
    accelcal->_conf_sample_time = sample_time;
    accelcal->_conf_fit_type = fit_type;

    const uint16_t faces = 2*accelcal->_conf_num_samples-4;
    const float a = (4.0f * M_PI / (3.0f * faces)) + M_PI / 3.0f;
    const float theta = 0.5f * acosf(cosf(a) / (1.0f - cosf(a)));
    accelcal->_min_sample_dist = GRAVITY_MSS * 2*sinf(theta/2);

    accelcal->_param.s.offset = *offset;
    accelcal->_param.s.diag = *diag;
    accelcal->_param.s.offdiag = *offdiag;

    switch (accelcal->_conf_fit_type) {
        case ACCEL_CAL_AXIS_ALIGNED_ELLIPSOID:
            if (accelcal->_conf_num_samples < 6) {
                accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
                return;
            }
            break;
        case ACCEL_CAL_ELLIPSOID:
            if (accelcal->_conf_num_samples < 8) {
                accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
                return;
            }
            break;
    }

    accelcal_set_status(accelcal, ACCEL_CAL_WAITING_FOR_ORIENTATION);
}

// returns true if accel calibrator is running
bool accelcal_running(accelcal_t accelcal)
{
    return accelcal->_status == ACCEL_CAL_WAITING_FOR_ORIENTATION || accelcal->_status == ACCEL_CAL_COLLECTING_SAMPLE;
}

// set Accel calibrator to start collecting samples in the next cycle
void accelcal_collect_sample(accelcal_t accelcal) {
    accelcal_set_status(accelcal, ACCEL_CAL_COLLECTING_SAMPLE);
}

// collect and avg sample to be passed onto LSQ estimator after all requisite orientations are done
void accelcal_new_sample(accelcal_t accelcal, const Vector3f_t* delta_velocity, float dt) {
    if (accelcal->_status != ACCEL_CAL_COLLECTING_SAMPLE) {
        return;
    }

    if (accelcal->_samples_collected >= accelcal->_conf_num_samples) {
        accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
        return;
    }

    //accelcal->_sample_buffer[accelcal->_samples_collected].delta_velocity += delta_velocity;
    vec3_add(&accelcal->_sample_buffer[accelcal->_samples_collected].delta_velocity,
             &accelcal->_sample_buffer[accelcal->_samples_collected].delta_velocity,
             delta_velocity);

    accelcal->_sample_buffer[accelcal->_samples_collected].delta_time += dt;

    accelcal->_last_samp_frag_collected_ms = time_millis();

    if (accelcal->_sample_buffer[accelcal->_samples_collected].delta_time > accelcal->_conf_sample_time) {
        Vector3f_t sample = vec3_mult2(&accelcal->_sample_buffer[accelcal->_samples_collected].delta_velocity, 1.0f/accelcal->_sample_buffer[accelcal->_samples_collected].delta_time);
        if (!accelcal_accept_sample(accelcal, &sample)) {
            accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
            return;
        }

        accelcal->_samples_collected++;

        if (accelcal->_samples_collected >= accelcal->_conf_num_samples) {
            accelcal_run_fit(accelcal, MAX_ITERATIONS, &accelcal->_fitness);

            if (accelcal->_fitness < accelcal->_conf_tolerance && accelcal_accept_result(accelcal)) {
                accelcal_set_status(accelcal, ACCEL_CAL_SUCCESS);
            } else {
                accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
            }
        } else {
            accelcal_set_status(accelcal, ACCEL_CAL_WAITING_FOR_ORIENTATION);
        }
    }
}

// interface for LSq estimator to read sample buffer sent after conversion from delta velocity
// to averaged acc over time
bool accelcal_get_sample(accelcal_t accelcal, uint8_t i, Vector3f_t* s) {
    if (i >= accelcal->_samples_collected) {
        return false;
    }
    vec3_mult(s, &accelcal->_sample_buffer[i].delta_velocity, 1.0f / accelcal->_sample_buffer[i].delta_time);
    return true;
}

// returns truen and sample corrected with diag offdiag parameters as calculated by LSq estimation procedure
// returns false if no correct parameter exists to be applied along with existing sample without corrections
bool accelcal_get_sample_corrected(accelcal_t accelcal, uint8_t i, Vector3f_t* s) {
    if (accelcal->_status != ACCEL_CAL_SUCCESS || !accelcal_get_sample(accelcal, i,s)) {
        return false;
    }

    matrix3f_t M = {
        {accelcal->_param.s.diag.x    , accelcal->_param.s.offdiag.x , accelcal->_param.s.offdiag.y},
        {accelcal->_param.s.offdiag.x , accelcal->_param.s.diag.y    , accelcal->_param.s.offdiag.z},
        {accelcal->_param.s.offdiag.y , accelcal->_param.s.offdiag.z , accelcal->_param.s.diag.z}
    };

    Vector3f_t sample_tmp = {s->x + accelcal->_param.s.offset.x,
                             s->y + accelcal->_param.s.offset.y,
                             s->z + accelcal->_param.s.offset.z};

    matrix3f_mul_vec(s, &M, &sample_tmp);

    return true;
}

// checks if no new sample has been received for considerable amount of time
void accelcal_check_for_timeout(accelcal_t accelcal) {
    const uint32_t timeout = accelcal->_conf_sample_time*2*1000 + 500;
    if (accelcal->_status == ACCEL_CAL_COLLECTING_SAMPLE && time_millis() - accelcal->_last_samp_frag_collected_ms > timeout) {
        accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
    }
}

// returns spherical fit paramteters
void accelcal_get_calibration(accelcal_t accelcal, Vector3f_t* offset) {
    offset->x = -accelcal->_param.s.offset.x;
    offset->y = -accelcal->_param.s.offset.y;
    offset->z = -accelcal->_param.s.offset.z;
}

// returns axis aligned ellipsoidal fit parameters
void accelcal_get_calibration2(accelcal_t accelcal, Vector3f_t* offset, Vector3f_t* diag) {
    offset->x = -accelcal->_param.s.offset.x;
    offset->y = -accelcal->_param.s.offset.y;
    offset->z = -accelcal->_param.s.offset.z;

    *diag = accelcal->_param.s.diag;
}

// returns generic ellipsoidal fit parameters
void accelcal_get_calibration3(accelcal_t accelcal, Vector3f_t* offset, Vector3f_t* diag, Vector3f_t* offdiag) {
    offset->x = -accelcal->_param.s.offset.x;
    offset->y = -accelcal->_param.s.offset.y;
    offset->z = -accelcal->_param.s.offset.z;

    *diag = accelcal->_param.s.diag;
    *offdiag = accelcal->_param.s.offdiag;
}

// sets status of calibrator and takes appropriate actions
static void accelcal_set_status(accelcal_t accelcal, enum accel_cal_status_t status)
{
    switch (status) {
        case ACCEL_CAL_NOT_STARTED:                 
            //Calibrator not started
            accelcal->_status = ACCEL_CAL_NOT_STARTED;

            accelcal->_samples_collected = 0;
            if (accelcal->_sample_buffer != NULL) {
                rt_free(accelcal->_sample_buffer);
                accelcal->_sample_buffer = NULL;
            }

            break;

        case ACCEL_CAL_WAITING_FOR_ORIENTATION:     
            //Callibrator has been started and is waiting for user to ack after orientation setting
            if (!accelcal_running(accelcal)) {
                accelcal->_samples_collected = 0;
                if (accelcal->_sample_buffer == NULL) {
                    accelcal->_sample_buffer = (struct AccelSample*)rt_calloc(accelcal->_conf_num_samples,sizeof(struct AccelSample));
                    if (accelcal->_sample_buffer == NULL) {
                        accelcal_set_status(accelcal, ACCEL_CAL_FAILED);
                        break;
                    }
                }
            }
            if (accelcal->_samples_collected >= accelcal->_conf_num_samples) {
                break;
            }
            accelcal->_status = ACCEL_CAL_WAITING_FOR_ORIENTATION;
            break;

        case ACCEL_CAL_COLLECTING_SAMPLE:
            // Calibrator is waiting on collecting samples from acceleromter for amount of 
            // time as requested by user/GCS
            if (accelcal->_status != ACCEL_CAL_WAITING_FOR_ORIENTATION) {
                break;
            }
            accelcal->_last_samp_frag_collected_ms = time_millis();
            accelcal->_status = ACCEL_CAL_COLLECTING_SAMPLE;
            break;

        case ACCEL_CAL_SUCCESS:
            // Calibrator has successfully fitted the samples to user requested surface model 
            // and has passed tolerance test
            if (accelcal->_status != ACCEL_CAL_COLLECTING_SAMPLE) {
                break;
            }

            accelcal->_status = ACCEL_CAL_SUCCESS;
            break;

        case ACCEL_CAL_FAILED:
            // Calibration has failed with reasons that can range from 
            // bad sample data leading to faillure in tolerance test to lack of distinct samples
            if (accelcal->_status == ACCEL_CAL_NOT_STARTED) {
                break;
            }

            accelcal->_status = ACCEL_CAL_FAILED;
            break;
    };
}

/*
  The sample acceptance distance is determined as follows:
  For any regular polyhedron with triangular faces, the angle theta subtended
  by two closest points is defined as
 
       theta = arccos(cos(A)/(1-cos(A)))
 
  Where:
       A = (4pi/F + pi)/3
  and
       F = 2V - 4 is the number of faces for the polyhedron in consideration,
       which depends on the number of vertices V
 
  The above equation was proved after solving for spherical triangular excess
  and related equations.
 */
static bool accelcal_accept_sample(accelcal_t accelcal, const Vector3f_t* sample)
{
    if (accelcal->_sample_buffer == NULL) {
        return false;
    }

    for(uint8_t i=0; i < accelcal->_samples_collected; i++) {
        Vector3f_t other_sample;
        accelcal_get_sample(accelcal, i, &other_sample);

        if (vec3_sub_length(&other_sample, sample) < accelcal->_min_sample_dist){
            return false;
        }
    }
    return true;
}

/*
    Run Gauss Newton fitting algorithm over the sample space and come up with offsets, diagonal/scale factors
    and crosstalk/offdiagonal parameters
*/
static void accelcal_run_fit(accelcal_t accelcal, uint8_t max_iterations, float* fitness)
{
    if (accelcal->_sample_buffer == NULL) {
        return;
    }
    *fitness = accelcal_calc_mean_squared_residuals(accelcal, &accelcal->_param.s);
    float min_fitness = *fitness;
    union accel_cal_param_u fit_param = accelcal->_param;
    uint8_t num_iterations = 0;

    while(num_iterations < max_iterations) {
        float JTJ[ACCEL_CAL_MAX_NUM_PARAMS*ACCEL_CAL_MAX_NUM_PARAMS];
        VectorP JTFI;

        rt_memset(JTJ, 0, ACCEL_CAL_MAX_NUM_PARAMS*ACCEL_CAL_MAX_NUM_PARAMS * sizeof(float));
        rt_memset(JTFI, 0, ACCEL_CAL_MAX_NUM_PARAMS * sizeof(float));

        for(uint16_t k = 0; k<accelcal->_samples_collected; k++) {
            Vector3f_t sample;
            accelcal_get_sample(accelcal, k, &sample);

            VectorP jacob;

            accelcal_calc_jacob(accelcal, &sample, &fit_param.s, jacob);

            for(uint8_t i = 0; i < accelcal_get_num_params(accelcal); i++) {
                // compute JTJ
                for(uint8_t j = 0; j < accelcal_get_num_params(accelcal); j++) {
                    JTJ[i*accelcal_get_num_params(accelcal)+j] += jacob[i] * jacob[j];
                }
                // compute JTFI
                JTFI[i] += jacob[i] * accelcal_calc_residual(&sample, &fit_param.s);
            }
        }

        if (!mat_inverse(JTJ, JTJ, accelcal_get_num_params(accelcal))) {
            return;
        }

        for(uint8_t row=0; row < accelcal_get_num_params(accelcal); row++) {
            for(uint8_t col=0; col < accelcal_get_num_params(accelcal); col++) {
                fit_param.a[row] -= JTFI[col] * JTJ[row*accelcal_get_num_params(accelcal)+col];
            }
        }

        *fitness = accelcal_calc_mean_squared_residuals(accelcal, &fit_param.s);

        if (isnan(*fitness) || isinf(*fitness)) {
            return;
        }

        if (*fitness < min_fitness) {
            min_fitness = *fitness;
            accelcal->_param = fit_param;
        }

        num_iterations++;
    }
}

// calculated the total mean squared fitness of all the collected samples using parameters
// supplied
static float accelcal_calc_mean_squared_residuals(accelcal_t accelcal, const struct accel_cal_param_t* params)
{
    if (accelcal->_sample_buffer == NULL || accelcal->_samples_collected == 0) {
        return 1.0e30f;
    }
    float sum = 0.0f;
    for(uint16_t i=0; i < accelcal->_samples_collected; i++){
        Vector3f_t sample;
        accelcal_get_sample(accelcal, i, &sample);
        float resid = accelcal_calc_residual(&sample, params);
        sum += sq(resid);
    }
    sum /= accelcal->_samples_collected;
    return sum;
}

// calculates residual from samples(corrected using supplied parameter) to gravity
// used to create Fitness column matrix
static float accelcal_calc_residual(const Vector3f_t* sample, const struct accel_cal_param_t* params) {
    matrix3f_t M = {
        {params->diag.x    , params->offdiag.x , params->offdiag.y},
        {params->offdiag.x , params->diag.y    , params->offdiag.z},
        {params->offdiag.y , params->offdiag.z , params->diag.z}
    };

    Vector3f_t tmp;
    Vector3f_t sample_tmp = {sample->x + params->offset.x,
                             sample->y + params->offset.y,
                             sample->z + params->offset.z};

    matrix3f_mul_vec(&tmp, &M, &sample_tmp);

    return GRAVITY_MSS - vec3_length(&tmp);
}

// calculate jacobian, a matrix that defines relation to variation in fitness with variation in each of the parameters
// this is used in LSq estimator to adjust variation in parameter to be used for next iteration of LSq
static void accelcal_calc_jacob(accelcal_t accelcal, const Vector3f_t* sample, const struct accel_cal_param_t* params, VectorP ret) {
    switch (accelcal->_conf_fit_type) {
        case ACCEL_CAL_AXIS_ALIGNED_ELLIPSOID:
        case ACCEL_CAL_ELLIPSOID: {
            const Vector3f_t* offset = &params->offset;
            const Vector3f_t* diag = &params->diag;
            const Vector3f_t* offdiag = &params->offdiag;
            matrix3f_t M = {
                {diag->x    , offdiag->x , offdiag->y},
                {offdiag->x , diag->y    , offdiag->z},
                {offdiag->y , offdiag->z , diag->z}
            };

            float A =  (diag->x    * (sample->x + offset->x)) + (offdiag->x * (sample->y + offset->y)) + (offdiag->y * (sample->z + offset->z));
            float B =  (offdiag->x * (sample->x + offset->x)) + (diag->y    * (sample->y + offset->y)) + (offdiag->z * (sample->z + offset->z));
            float C =  (offdiag->y * (sample->x + offset->x)) + (offdiag->z * (sample->y + offset->y)) + (diag->z    * (sample->z + offset->z));

            Vector3f_t tmp;
            Vector3f_t sample_tmp = {sample->x + offset->x,
                                     sample->y + offset->y,
                                     sample->z + offset->z};

            matrix3f_mul_vec(&tmp, &M, &sample_tmp);
            float length = vec3_length(&tmp);

            // 0-2: offsets
            ret[0] = -1.0f * (((diag->x    * A) + (offdiag->x * B) + (offdiag->y * C))/length);
            ret[1] = -1.0f * (((offdiag->x * A) + (diag->y    * B) + (offdiag->z * C))/length);
            ret[2] = -1.0f * (((offdiag->y * A) + (offdiag->z * B) + (diag->z    * C))/length);
            // 3-5: diagonals
            ret[3] = -1.0f * ((sample->x + offset->x) * A)/length;
            ret[4] = -1.0f * ((sample->y + offset->y) * B)/length;
            ret[5] = -1.0f * ((sample->z + offset->z) * C)/length;
            // 6-8: off-diagonals
            ret[6] = -1.0f * (((sample->y + offset->y) * A) + ((sample->x + offset->x) * B))/length;
            ret[7] = -1.0f * (((sample->z + offset->z) * A) + ((sample->x + offset->x) * C))/length;
            ret[8] = -1.0f * (((sample->z + offset->z) * B) + ((sample->y + offset->y) * C))/length;
            return;
        }
    }
}

// returns number of parameters are required for selected Fit type
static uint8_t accelcal_get_num_params(accelcal_t accelcal) {
    switch (accelcal->_conf_fit_type) {
        case ACCEL_CAL_AXIS_ALIGNED_ELLIPSOID:
            return 6;
        case ACCEL_CAL_ELLIPSOID:
            return 9;
        default:
            return 6;
    }
}

// determines if the result is acceptable
static bool accelcal_accept_result(accelcal_t accelcal) {
    if (fabsf(accelcal->_param.s.offset.x) > GRAVITY_MSS ||
        fabsf(accelcal->_param.s.offset.y) > GRAVITY_MSS ||
        fabsf(accelcal->_param.s.offset.z) > GRAVITY_MSS ||
        accelcal->_param.s.diag.x < 0.8f || accelcal->_param.s.diag.x > 1.2f ||
        accelcal->_param.s.diag.y < 0.8f || accelcal->_param.s.diag.y > 1.2f ||
        accelcal->_param.s.diag.z < 0.8f || accelcal->_param.s.diag.z > 1.2f) {
        return false;
    } else {
        return true;
    }
}

/*------------------------------------test------------------------------------*/


